#---------------------------------------------------------------
#							图像滤波
#---------------------------------------------------------------
from PIL import Image
from PIL import ImageFilter
import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

class image_filtering():
	def __init__(self, mode):
		self.mode = mode
		self.img_gray = None
		self.width = None
		self.height = None

	def write_txt(self, name, _list):
		with open (f'..\\image-data\\{name}.txt', 'w') as f:
			for i in range(len(_list)):
				if self.mode == 'sharpen':
					f.write(f'{_list[i]:08b}')
				else:
					temp = _list[i] - 128
					if temp < 0:
						temp = temp + 256
					f.write(f'{temp:08b}')
				if i != len(_list)-1:
					f.write('\n')

	def load_img_gray(self, img_gray):
		self.img_gray = img_gray
		f = Image.open(f'..\\image\\{img_gray}.tif')
		pix = f.load()
		self.width = f.size[0]
		self.height = f.size[1]
		pix_data = []
		for i in range(self.height):
			for j in range(self.width):
				pix_data.append(pix[j,i])
		self.write_txt(f'{img_gray}_gray', pix_data)
		f.show()
		#f.save(f'..\\image-processed\\{self.img_gray}_gray_origin.tif')			

	def img_data_reshape(self, img_data):
		_list = [[] for i in range(self.height)]
		for i in range(self.width):
			for j in range(self.height):
				_list[j].append(int(img_data[i+j*self.width]))
		return _list

	def load_img_data(self, name):
		with open(f'..\\image-data\\{name}.txt', 'r', newline = '') as f:
			read_stream = f.read().splitlines()
		data_reshape = self.img_data_reshape(read_stream)
		data_array = np.array(data_reshape)
		return data_array

	def load_img_gray_data(self, mode):
		pix_data = self.load_img_data(f'{self.img_gray}_gray_{mode}')
		img_data = [[] for i in range(self.height)]
		for i in range(self.height):
			for j in range(self.width):
				img_data[i].append(pix_data[i, j])	
		return img_data

	def load_img_gray_data_sharpen(self):
		img_data = self.load_img_gray_data(f'sharpen')
		img = Image.fromarray(np.uint8(img_data))
		img.show()
		img.save(f'..\\image-processed\\{self.img_gray}_gray_sharpen.tif')


	def load_img_gray_data_edge_detect(self):
		img_data_x = self.load_img_gray_data(f'edge_detect_x')
		img_data_y = self.load_img_gray_data(f'edge_detect_y')
		img_edge_detect = [[] for i in range(self.height)]
		for i in range(len(img_data_x)):
			for j in range(len(img_data_x[i])):
				if img_data_x[i][j] < 0:
					temp_x = 0
				else:
					temp_x = img_data_x[i][j]

				if img_data_y[i][j] < 0:
					temp_y = 0
				else:
					temp_y = img_data_y[i][j]	


				if round(math.sqrt(temp_x**2 + temp_y**2)) > 255:
					temp = 255
				else:
					temp = round(math.sqrt(temp_x**2 + temp_y**2))
				img_edge_detect[i].append(temp)
		img = Image.fromarray(np.uint8(img_edge_detect))
		img.show()
		img.save(f'..\\image-processed\\{self.img_gray}_gray_edge_detect.tif')

	def load_img_gray_filtering_data(self):
		if self.mode == f'sharpen':
			self.load_img_gray_data_sharpen()
		elif self.mode == f'edge_detect':
			self.load_img_gray_data_edge_detect()
		else:
			print("Error! Filter type wrong.")
		psnr = self.calculate_psnr(f'{self.img_gray}\\{self.img_gray}_gray_{self.mode}', f'{self.img_gray}_gray_{self.mode}')
		ssim = self.calculate_ssim(f'{self.img_gray}\\{self.img_gray}_gray_{self.mode}', f'{self.img_gray}_gray_{self.mode}')
		print(f'PSNR:{psnr}, SSIM:{ssim}')

	def calculate_psnr(self, img_accurate, img_approximate):
		#load accurate img
		f = Image.open(f'..\\image-processed\\{img_accurate}.tif')
		pix = f.load()
		width = f.size[0]
		height = f.size[1]
		img_accurate_data = []
		for i in range(height):
			for j in range(width):
				img_accurate_data.append(pix[j,i])

		#load approximate img
		f = Image.open(f'..\\image-processed\\{img_approximate}.tif')
		pix = f.load()
		img_approximate_data = []
		for i in range(height):
			for j in range(width):
				img_approximate_data.append(pix[j,i])

		img_accurate_data = np.array(img_accurate_data)
		img_approximate_data = np.array(img_approximate_data)

		return 20 * math.log10(255/(math.sqrt(np.mean((img_approximate_data - img_accurate_data)**2))))


	def calculate_ssim(self, img_accurate, img_approximate):
		#load accurate img
		f = Image.open(f'..\\image-processed\\{img_accurate}.tif')
		pix = f.load()
		width = f.size[0]
		height = f.size[1]
		img_accurate_data = []
		for i in range(height):
			for j in range(width):
				img_accurate_data.append(pix[j,i])

		#load approximate img
		f = Image.open(f'..\\image-processed\\{img_approximate}.tif')
		pix = f.load()
		img_approximate_data = []
		for i in range(height):
			for j in range(width):
				img_approximate_data.append(pix[j,i])

		img_accurate_data = np.array(img_accurate_data)
		img_approximate_data = np.array(img_approximate_data)

		mean_img_accurate_data = np.mean(img_accurate_data)
		mean_img_approximate_data = np.mean(img_approximate_data)
		variance_img_accurate_data = math.sqrt(sum((img_accurate_data - mean_img_accurate_data)**2) / (len(img_accurate_data) - 1))
		variance_img_approximate_data = math.sqrt(sum((img_approximate_data - mean_img_approximate_data)**2) / (len(img_accurate_data) - 1))
		covariance_img_data = sum((img_accurate_data - mean_img_accurate_data) * (img_approximate_data - mean_img_approximate_data)) / (len(img_accurate_data) - 1)

		k1 = 0.01
		k2 = 0.03
		l = 255
		c1 = (k1*l)**2
		c2 = (k2*l)**2
		c3 = c2/2

		ssim_l = (2*mean_img_accurate_data*mean_img_approximate_data + c1) / (mean_img_accurate_data**2 + mean_img_approximate_data**2 + c1)
		ssim_c = (2*variance_img_accurate_data*variance_img_approximate_data + c2) / (variance_img_accurate_data**2 + variance_img_approximate_data**2 + c2)
		ssim_s = (covariance_img_data + c3) / (variance_img_accurate_data*variance_img_approximate_data + c3)

		return ssim_l*ssim_c*ssim_s

#--------------------------------
# Mode is sharpen or edge_detect
#--------------------------------

img_filter = image_filtering(f'edge_detect')
img_filter.load_img_gray('Peppers')
img_filter.load_img_gray_filtering_data()
